﻿using System;
using System.Collections.Generic;

using System.IO;
using System.Linq;
using System.Runtime.Serialization.Formatters.Binary;
using System.Text;
using System.Xml;

using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;

namespace RayDen.Library.Components.PhotonMapping {
    public class KdTreePhotonMap : IPhotonMap {
        private Photon[] photons;
        int max_photons;
        internal int photonsCount;
        int prev_scale;

        public AABB bbox;

        public ulong PhotonsCount { get { return (ulong)this.photonsCount; } }
        

        public Photon[] GetPlainData() {
            return this.photons.ToArray();
        }

        public KdTreePhotonMap(int max) {
            this.max_photons = max;
            this.prev_scale = 1;
            this.photonsCount = 0;
            this.bbox = new AABB();
            this.photons = new Photon[max];
        }


        public static Vector PhotonDirection(ref Photon p) {
            return new Vector(
                DirectionDiscretizationHelper.SinTheta(p.Theta) * DirectionDiscretizationHelper.CosPhi(p.Phi),
                DirectionDiscretizationHelper.SinTheta(p.Theta) * DirectionDiscretizationHelper.SinPhi(p.Phi),
                DirectionDiscretizationHelper.CosTheta(p.Theta)
                );
        }
        public void QueryMap(ref Point pos, float maxDist, int nphotons, out Photon[] result) {
            var np = new NearestPhotons();
            np.dist = new float[nphotons + 1];
            np.index = Enumerable.Repeat(new Photon() {Plane = 255}, nphotons + 1).ToArray();
            np.pos = pos;
            np.max = nphotons;
            np.gotHeap = 0;
            np.found = 0;

            np.dist[0] = maxDist * maxDist;
            this.LocatePhotons(ref np, 1);
            result = np.found > 0 ? np.index.Skip(1).ToArray() : new Photon[0];
        }

        public Vector IrradianceEstimate(ref Point pos, ref Normal n, float max_dist, int nphotons) {
            var normal = n;
            var irrad = new RgbSpectrum(0.0f);
            var np = new NearestPhotons();
            np.dist = new float[nphotons + 1];
            np.index = new Photon[nphotons + 1];
            np.pos = pos;
            np.max = nphotons;
            np.gotHeap = 0;
            np.found = 0;

            np.dist[0] = max_dist * max_dist;

            this.LocatePhotons(ref np, 1);

            if (np.found < 1)
                return new Vector(0);

            for (int index = 0; index < np.index.Length; index++) {
                Vector pdir = 
                    PhotonDirection(ref np.index[index]);
                var dot = Vector.Dot(ref pdir, ref normal);

                 if (dot > 0.0f) 
                {
                    irrad = irrad + np.index[index].Power; //*BRDF
                }
            }

            var tmp = (1.0f / MathLab.M_PI) / np.dist[0];

            irrad = irrad * tmp;

            return irrad;

        }

        public int Half;
        private void LocatePhotonsTemp(ref NearestPhotons np, int index) {
            Photon p = this.photons[index];
            float dist1;

            if (index < this.Half) {
                dist1 = np.pos[p.Plane] - p.Position[p.Plane];

                if (dist1 > 0.0d) {
                    this.LocatePhotons(ref np, 2 * index + 1);

                    if ((dist1 * dist1) < np.dist[0])
                        this.LocatePhotons(ref np, 2 * index);
                }
                else {
                    this.LocatePhotons(ref np, 2 * index);
                    if ((dist1 * dist1) < np.dist[0])
                        this.LocatePhotons(ref np, 2 * index + 1);
                }
            }

            dist1 = (p.Position - (Vector)np.pos).Length;
            float dist2 = dist1 * dist1;

            if (dist2 < np.dist[0]) {
                //Photon found
                if (np.found < np.max) {
                    np.found++;
                    np.dist[np.found] = dist2;
                    np.index[np.found] = p;
                }
                else {
                    int j, parent;

                    if (np.gotHeap == 0) {
                        double dst2;
                        Photon phot;
                        int hFound = np.found >> 1;
                        for (int k = hFound; k >= 1; k--) {
                            parent = k;
                            phot = np.index[k];
                            dst2 = np.dist[k];

                            while (parent <= hFound) {
                                j = parent + parent;
                                if (j < np.found && np.dist[j] < np.dist[j + 1]) {
                                    j++;
                                }
                                if (dst2 >= np.dist[j])
                                    break;
                                np.dist[parent] = np.dist[j];
                                np.index[parent] = phot;
                                parent = j;
                            }
                        }
                        np.gotHeap = 1;
                    }

                    parent = 1;
                    j = 2;
                    while (j <= np.found) {
                        if (j < np.found && np.dist[j] < np.dist[j + 1])
                            j++;
                        if (dist2 > np.dist[j])
                            break;
                        np.dist[parent] = np.dist[j];
                        np.index[parent] = np.index[j];
                        parent = j;
                        j += j;
                    }
                    np.index[parent] = p;
                    np.dist[parent] = dist2;

                    np.dist[0] = np.dist[1];
                }
            }
        }

        private void LocatePhotons(ref NearestPhotons np, int index) {
            float dist1;

            if (index < this.Half) {
                dist1 = np.pos[this.photons[index].Plane] - this.photons[index].Position[this.photons[index].Plane];

                if (dist1 > 0.0d) {
                    this.LocatePhotons(ref np, 2 * index + 1);

                    if ((dist1 * dist1) < np.dist[0])
                        this.LocatePhotons(ref np, 2 * index);
                }
                else {
                    this.LocatePhotons(ref np, 2 * index);
                    if ((dist1 * dist1) < np.dist[0])
                        this.LocatePhotons(ref np, 2 * index + 1);
                }
            }

            dist1 = Geometry.EuclidianDistance(ref this.photons[index].Position, ref np.pos);
            /*
                (new Vector( this.photons[index].Position.x - np.pos.x, 
                                 this.photons[index].Position.y - np.pos.y,
                                 this.photons[index].Position.z - np.pos.z)).Length;
            */
            float dist2 = dist1 * dist1;

            if (dist2 < np.dist[0]) {
                //Photon found
                if (np.found < np.max) {
                    np.found++;
                    np.dist[np.found] = dist2;
                    np.index[np.found] = this.photons[index];
                }
                else {
                    int j, parent;

                    if (np.gotHeap == 0) {
                        double dst2;
                        Photon phot;
                        int hFound = np.found >> 1;
                        for (int k = hFound; k >= 1; k--) {
                            parent = k;
                            phot = np.index[k];
                            dst2 = np.dist[k];

                            while (parent <= hFound) {
                                j = parent + parent;
                                if (j < np.found && np.dist[j] < np.dist[j + 1]) {
                                    j++;
                                }
                                if (dst2 >= np.dist[j])
                                    break;
                                np.dist[parent] = np.dist[j];
                                np.index[parent] = phot;
                                parent = j;
                            }
                        }
                        np.gotHeap = 1;
                    }

                    parent = 1;
                    j = 2;
                    while (j <= np.found) {
                        if (j < np.found && np.dist[j] < np.dist[j + 1])
                            j++;
                        if (dist2 > np.dist[j])
                            break;
                        np.dist[parent] = np.dist[j];
                        np.index[parent] = np.index[j];
                        parent = j;
                        j += j;
                    }
                    np.index[parent] = this.photons[index];
                    np.dist[parent] = dist2;

                    np.dist[0] = np.dist[1];
                }
            }
        }

        public void Store(ref RgbSpectrum power, ref Point pos, ref Vector dir) {
            if (this.photonsCount > (this.max_photons - 2))
                return;

            if (power.Length <= 0f)
                return;

            Photon node = this.photons[this.photonsCount];
            //node.Direction = dir.Normalize();
            node.Power = power;
            node.Position = pos;
            this.bbox = this.bbox.Union(node.Position);
            //if (node.position < bbox.BBoxMin)
            //    bbox.BBoxMin = node.position;

            //if (node.position > bbox.BBoxMax)
            //    bbox.BBoxMax = node.position;
            var dirx = MathLab.Clamp(dir.z, -1f, 1f);
            int theta = Convert.ToInt32(Math.Acos(dirx) * (256.0d / Math.PI));
            if (theta > 255)
                theta = 255;

            node.Theta = (byte)theta;

            int phi = Convert.ToInt32(Math.Atan2(dir.y, dir.x) * (256.0d / 2.0 * Math.PI));


            if (phi > 255)
                phi = 255;
            else

                if (phi < 0)
                    phi = phi + 256;

            node.Phi = (byte)phi;
            this.photons[this.photonsCount] = node;
            this.photonsCount++;

        }

        public void ScalePhotonMap(float scale) {
            for (int i = this.prev_scale; i < this.photonsCount; i++) {
                this.photons[i].SetPower(this.photons[i].Power * scale);
            }

            this.prev_scale = this.photonsCount;
        }

        public void Balance() {

            /*
              if (storedPhotons == 0)
            return;
        photons = photonList.toArray(new Photon[photonList.size()]);
        photonList = null;
        Photon[] temp = new Photon[storedPhotons + 1];
        balanceSegment(temp, 1, 1, storedPhotons);
        photons = temp;
        halfStoredPhotons = storedPhotons / 2;
        log2n = (int) Math.ceil(Math.log(storedPhotons) / Math.log(2.0));

             
             
             */
            if (this.photonsCount == 0)
                return;

            Photon[] photonz = new Photon[this.photons.Length];
            Array.Copy(this.photons.ToArray(), photonz, this.photons.Length);
            List<Photon> photonList = new List<Photon>(photonz);
            photonz = null;

            var temp = Enumerable.Repeat(new Photon(), photonsCount + 1).ToList();

            this.BalanceSegment(ref temp, ref photonList, 1, 1, this.photonsCount);

            this.photons = temp.ToArray();
            this.Half = this.photons.Length / 2;

            //if (photonsCount > 1)
            //{
            //    List<Photon> pa1 = new List<Photon>(photonsCount + 1);
            //    List<Photon> pa2 = new List<Photon>(photonsCount + 1);
            //    PrepareList(photonsCount + 1,ref pa1);
            //    PrepareList(photonsCount + 1,ref pa2);

            //    for (int i =0 ;i <photonsCount; i++)
            //    {
            //        pa2[i] = photons[i];
            //    }

            //    BalanceSegment(ref pa1, ref pa2, 1, 1, photonsCount);
            //    pa2 = null;
            //    int d;
            //    int j = 1;
            //    int foo = 1;

            //    Photon fooPhoton = photons[j];
            //    for(int i = 0; i < photonsCount; i++)
            //    {
            //        d = pa1.IndexOf(pa1[j])  - photonsCount; // TODO INCORRECT
            //        pa1[j] = Photon.NullPhoton;
            //       // pa1.Remove(pa1[j]);

            //        if (d != foo)
            //            photons[j] = photons[d];
            //        else
            //        {
            //            photons[j] = fooPhoton;

            //            if (i < photonsCount)
            //            {
            //                for (; foo <= photonsCount; foo ++ )
            //                {
            //                    if ((!Photon.IsNull(pa1[foo])))
            //                        break;
            //                    fooPhoton = photons[foo];
            //                    j = foo;
            //                }
            //                continue;
            //            }
            //            j = d;
            //        }
            //    }
            //    pa1 = null;

            //    Half = photonsCount/2 + 1;

            //}
        }

        private void Swap(ref List<Photon> arr, int a, int b) {
            Photon tmp = arr[a];
            arr[a] = arr[b];
            arr[b] = tmp;
        }

        private void MedianSplit(ref List<Photon> p, int start, int end, int median, short axis) {
            int left = start;
            int right = end;

            while (right > left) {
                double v = p[right].Position[axis];
                int i = left - 1;
                int j = right;
                for (; ; ) {
                    while (p[++i].Position[axis] < v)
                        ;
                    while (p[--j].Position[axis] > v && j > left)
                        ;
                    if (i >= j)
                        break;
                    this.Swap(ref p, i, j);
                }

                this.Swap(ref p, i, right);
                if (i >= median)
                    right = i - 1;
                if (i <= median)
                    left = i + 1;
            }
        }


        private void BalanceSegment(ref List<Photon> pbal, ref List<Photon> porg, int index, int start, int end) {
            int median = 1;
            while ((4 * median) <= (end - start + 1))
                median += median;
            if ((3 * median) <= (end - start + 1)) {
                median += median;
                median += start - 1;
            }
            else {
                median = end - median + 1;
            }

            byte axis = 2;
            Vector dist = this.bbox.Max - this.bbox.Min;
            if ((dist.x > dist.y) && (dist.x > dist.z))
                axis = 0;
            else
                if (dist.y > dist.z)
                    axis = 1;
            this.MedianSplit(ref porg, start, end, median, axis);
            pbal[index] = porg[median];
            pbal[index].SetPlane(axis);

            if (median > start) {
                if (start < median - 1) {
                    var tmp = this.bbox.Max[axis];
                    this.bbox.Max.SetComponent(axis,pbal[index].Position[axis] );
                    this.BalanceSegment(ref pbal, ref porg, 2 * index, start, median - 1);
                    this.bbox.Max.SetComponent(axis,tmp);
                }
                else {
                    pbal[2 * index] = porg[start];
                }
            }

            if (median < end) {
                if (end > median + 1) {
                    var tmp = this.bbox.Min[axis];
                    this.bbox.Min.SetComponent(axis,pbal[index].Position[axis]);
                    this.BalanceSegment(ref pbal, ref porg, 2 * index + 1, median + 1, end);
                    this.bbox.Min.SetComponent(axis,tmp);
                }
                else {
                    pbal[2 * index + 1] = porg[end];
                }
            }
        }

        public void SavePhotonMap(string fileName) {
            BinaryFormatter frm = new BinaryFormatter();
            using (FileStream fs = new FileStream(fileName, FileMode.Create, FileAccess.ReadWrite)) {
                BinaryWriter writer = new BinaryWriter(fs);
                string xmlInfo = string.Format(@"<PhotonMapInfo> <photonsCount>{0}</photonsCount> <maxPhotons>{1}</maxPhotons><Half>{2}</Half><PrevScale>{3}</PrevScale>{4}</PhotonMapInfo>",
                                               this.photonsCount, this.max_photons, this.Half, this.prev_scale, this.bbox.ToString());
                byte len = (byte)xmlInfo.Length;
                fs.WriteByte(len);
                fs.Write(Encoding.UTF8.GetBytes(xmlInfo), 0, xmlInfo.Length);
                foreach (Photon photon in this.photons) {
                    writer.Write(photon.Position.x);
                    writer.Write(photon.Position.y);
                    writer.Write(photon.Position.z);
                    writer.Write(photon.Power.c1);
                    writer.Write(photon.Power.c2);
                    writer.Write(photon.Power.c3);
                    writer.Write(photon.Phi);
                    writer.Write(photon.Theta);
                    writer.Write(photon.Plane);
//                    frm.Serialize(fs, photon);
                }
                fs.Close();
                writer.Dispose();
            }
        }


        public void LoadPhotonMap(string fileName) {
            BinaryFormatter frm = new BinaryFormatter();
            using (FileStream fs = new FileStream(fileName, FileMode.Open, FileAccess.Read)) {
                string xmlInfo = string.Empty;
                int len = fs.ReadByte();
                byte[] xml = new byte[len];
                fs.Read(xml, 0, len);
                xmlInfo = new string(Encoding.UTF8.GetChars(xml));

                XmlDocument param = new XmlDocument();
                param.LoadXml(xmlInfo);
                foreach (XmlNode node in param.ChildNodes[0].ChildNodes) {
                    if (node.Name == "photonsCount") {
                        this.photonsCount = int.Parse(node.InnerText.ToString());
                    }

                    if (node.Name == "maxPhotons") {
                        this.max_photons = int.Parse(node.InnerText.ToString());
                    }
                    if (node.Name == "Half") {
                        this.Half = int.Parse(node.InnerText.ToString());
                    }
                    if (node.Name == "PrevScale") {
                        this.prev_scale = int.Parse(node.InnerText.ToString());
                    }
                    if (node.Name == "BBox") {
                        string nodeVal = node.InnerText.ToString();
                    }
                }

                this.photons = new Photon[max_photons];
                var sp = fs.Position;
                using (BinaryReader reader = new BinaryReader(fs))
                    for (int index = 0; index < this.photons.Length; index++) {
                        var position = new Point(reader.ReadSingle(), reader.ReadSingle(), reader.ReadSingle());
                        var power = new RgbSpectrum(reader.ReadSingle(), reader.ReadSingle(), reader.ReadSingle());
                        byte phi = reader.ReadByte();
                        byte theta = reader.ReadByte();
                        byte plane = reader.ReadByte();
                        this.photons[index] = new Photon() {Position =  position, Power =  power, Phi = phi, Theta = theta, Plane = plane};
                    }

                fs.Close();
            }
        }
    }
}